# db_service.py
import asyncio
import logging
import os
import sqlite3
from pathlib import Path
from typing import Any, Dict, List, Optional, Union, AsyncGenerator
from contextlib import asynccontextmanager
from pydantic import BaseModel, Field
import aiosqlite


# 自定义异常
class DatabaseError(Exception):
    """数据库操作基础异常"""

    def __init__(self, message: str, query: Optional[str] = None):
        super().__init__(message)
        self.query = query


class TableExistsError(DatabaseError):
    """表已存在异常"""
    pass


class RecordNotFoundError(DatabaseError):
    """记录未找到异常"""
    pass


# 数据模型
class DatabaseConfig(BaseModel):
    """数据库配置项"""
    db_path: Path = Path(os.path.join(os.getcwd() , 'data', 'data.db'))
    timeout: float = 10.0
    journal_mode: str = "WAL"  # 使用Write-Ahead Logging模式提升性能
    busy_timeout: int = 5000  # 单位毫秒


class TableSchema(BaseModel):
    """表结构定义模型"""
    name: str
    columns: Dict[str, str]  # 列名: 类型
    indexes: List[str] = Field(default_factory=list)


class QueryCondition(BaseModel):
    """查询条件模型"""
    where: Optional[str] = None
    params: Union[tuple, Dict[str, Any], None] = None
    order_by: Optional[str] = None
    limit: Optional[int] = None
    offset: Optional[int] = None


class SqlLiteService:
    """异步SQLite数据库服务"""

    def __init__(self, config: DatabaseConfig, logger: Optional[logging.Logger] = None):
        self.config = config
        self.logger = logger or logging.getLogger(__name__)
        self._conn: Optional[aiosqlite.Connection] = None
        self._is_initialized = False

        # 预定义表结构
        self.table_schemas = {
            "workstations": TableSchema(
                name="workstations",
                columns={
                    "id": "INTEGER PRIMARY KEY AUTOINCREMENT",
                    "workstation_id": "TEXT UNIQUE NOT NULL",
                    "name": "TEXT NOT NULL",
                    "config": "TEXT",
                    "is_open": "BOOLEAN DEFAULT 0",
                    "created_at": "TIMESTAMP DEFAULT CURRENT_TIMESTAMP",
                    "last_activity": "TIMESTAMP"
                },
                indexes=[
                    "idx_workstation_id ON workstations(workstation_id)",
                    "idx_is_open ON workstations(is_open)"
                ]
            )
        }

    @asynccontextmanager
    async def get_connection(self) -> AsyncGenerator[aiosqlite.Connection, None]:
        """获取数据库连接上下文管理器"""
        try:
            if not self._conn:
                self._conn = await aiosqlite.connect(
                    self.config.db_path,
                    timeout=self.config.timeout
                )
                # 优化配置
                await self._conn.execute(f"PRAGMA journal_mode={self.config.journal_mode}")
                await self._conn.execute(f"PRAGMA busy_timeout={self.config.busy_timeout}")
                self.logger.info(f"Connected to database: {self.config.db_path}")

            yield self._conn
        except Exception as e:
            self.logger.error(f"Database connection failed: {str(e)}, db_path: {self.config.db_path}")
            raise DatabaseError(f"Connection error: {str(e)}") from e

    async def initialize(self):
        """初始化数据库表结构"""
        if self._is_initialized:
            return

        async with self.get_connection() as conn:
            try:
                # 创建表
                for table_name, schema in self.table_schemas.items():
                    if await self._table_exists(conn, table_name):
                        continue

                    columns = ", ".join(
                        [f"{col_name} {col_type}"
                         for col_name, col_type in schema.columns.items()]
                    )
                    create_table_sql = f"CREATE TABLE {table_name} ({columns})"

                    await conn.execute(create_table_sql)

                    # 创建索引
                    for index_sql in schema.indexes:
                        await conn.execute(f"CREATE INDEX {index_sql}")

                await conn.commit()
                self._is_initialized = True
                self.logger.info("Database tables initialized")

            except aiosqlite.OperationalError as e:
                if "already exists" in str(e):
                    raise TableExistsError(f"Table already exists: {table_name}") from e
                raise DatabaseError(f"Initialization failed: {str(e)}") from e

    async def _table_exists(self, conn: aiosqlite.Connection, table_name: str) -> bool:
        """检查表是否存在"""
        query = """
            SELECT count(*) FROM sqlite_master 
            WHERE type='table' AND name=?
        """
        cursor = await conn.execute(query, (table_name,))
        result = await cursor.fetchone()
        return result[0] > 0 if result else False

    async def select_data(
            self,
            table: str,
            condition: Union[str, QueryCondition, None] = None,
            params: Union[tuple, Dict[str, Any], None] = None,
            single: bool = False
    ) -> Union[Dict, List[Dict], None]:
        """
        通用查询方法
        :param single: 是否只返回单条记录
        """
        condition = QueryCondition(where=condition, params=params)


        query = f"SELECT * FROM {table}"
        params = params or ()

        if condition.where:
            query += f" WHERE {condition.where}"
            params = condition.params or ()

        if condition.order_by:
            query += f" ORDER BY {condition.order_by}"

        if condition.limit is not None:
            query += f" LIMIT {condition.limit}"
            if condition.offset is not None:
                query += f" OFFSET {condition.offset}"

        self.logger.debug(f"Query: {query}")
        self.logger.debug(f"params: {params}")
        async with self.get_connection() as conn:
            conn.row_factory = aiosqlite.Row  # 返回字典形式结果
            cursor = await conn.execute(query, params)
            results = await cursor.fetchall()
            self.logger.debug(f"Results: {results}")

            if not results:
                if single:
                    return None
                return None if single else []

            return [dict(row) for row in results] if not single else dict(results[0])

    async def insert_data(
            self,
            table: str,
            data: Dict[str, Any],
            on_conflict: str = "IGNORE"
    ) -> int:
        """
        插入数据
        :return: 插入的行ID
        """
        columns = ", ".join(data.keys())
        placeholders = ", ".join("?" * len(data))
        values = tuple(data.values())

        query = f"""
            INSERT OR {on_conflict} INTO {table} ({columns})
            VALUES ({placeholders})
        """

        async with self.get_connection() as conn:
            try:
                cursor = await conn.execute(query, values)
                await conn.commit()
                return cursor.lastrowid
            except aiosqlite.IntegrityError as e:
                self.logger.error(f"Insert failed: {str(e)}")
                raise DatabaseError(f"Insert conflict: {str(e)}", query) from e

    async def update_data(
            self,
            table: str,
            data: Dict[str, Any],  # 调整参数顺序
            where: str,
            params: Optional[tuple] = None
    ) -> int:
        """
        更新数据
        :return: 影响的行数
        """
        set_clause = ", ".join([f"{k} = ?" for k in data.keys()])
        values = tuple(data.values()) + (params or ())

        query = f"""
            UPDATE {table} 
            SET {set_clause}
            WHERE {where}
        """

        async with self.get_connection() as conn:
            cursor = await conn.execute(query, values)
            await conn.commit()
            return cursor.rowcount

    async def delete_data(
            self,
            table: str,
            where: str,
            params: Optional[tuple] = None
    ) -> int:
        """删除数据"""
        query = f"DELETE FROM {table} WHERE {where}"

        async with self.get_connection() as conn:
            cursor = await conn.execute(query, params or ())
            await conn.commit()
            return cursor.rowcount

    async def execute_raw_sql(
            self,
            sql: str,
            params: Optional[tuple] = None,
            commit: bool = True
    ) -> aiosqlite.Cursor:
        """执行原始SQL语句"""
        async with self.get_connection() as conn:
            cursor = await conn.execute(sql, params or ())
            if commit:
                await conn.commit()
            return cursor

    @asynccontextmanager
    async def transaction(self):
        """事务上下文管理器"""
        async with self.get_connection() as conn:
            try:
                await conn.execute("BEGIN")
                yield
                await conn.commit()
            except Exception as e:
                await conn.rollback()
                self.logger.error(f"Transaction rolled back: {str(e)}")
                raise DatabaseError(f"Transaction failed: {str(e)}") from e

    async def close(self):
        """关闭数据库连接"""
        if self._conn:
            await self._conn.close()
            self._conn = None
            self.logger.info("Database connection closed")


# 使用示例
async def main():
    # 初始化配置
    config = DatabaseConfig(db_path=Path("test.db"))

    # 创建服务实例
    db_service = SqlLiteService(config)

    try:
        # 初始化数据库
        await db_service.initialize()

        # 插入数据
        workstation_id = await db_service.insert_data("workstations", {
            "workstation_id": "ws_001",
            "name": "Test Workstation",
            "is_open": 1
        })

        # 查询数据
        records = await db_service.select_data(
            "workstations",
            QueryCondition(
                where="is_open = ?",
                params=(1,),
                order_by="created_at DESC",
                limit=10
            )
        )
        print(records)

    finally:
        await db_service.close()
